# -*- coding: utf-8 -*-
import os, json, joblib
from typing import Dict, Optional
from dataclasses import dataclass
import numpy as np

MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
META_PATH  = os.path.join(MODELS_DIR, "_active.json")

SUPPORTED = {
    "crop":  [".pkl", ".joblib"],   # 作物推荐
    "yield": [".pkl", ".joblib"],   # 产量预测
}

@dataclass
class ModelPack:
    kind: str
    path: str
    obj: object

class ModelRegistry:
    def __init__(self):
        os.makedirs(MODELS_DIR, exist_ok=True)
        self.active: Dict[str, Optional[str]] = {"crop": None, "yield": None}
        self.loaded: Dict[str, ModelPack] = {}
        self._load_meta()

    def _load_meta(self):
        if os.path.exists(META_PATH):
            self.active.update(json.load(open(META_PATH, "r", encoding="utf-8")))

    def _save_meta(self):
        json.dump(self.active, open(META_PATH, "w", encoding="utf-8"),
                  ensure_ascii=False, indent=2)

    def list_active(self): return self.active

    def _check_ext(self, kind: str, filename: str):
        ok = os.path.splitext(filename.lower())[1] in SUPPORTED[kind]
        if not ok: raise ValueError(f"{kind} 只支持 {SUPPORTED[kind]}")

    def save_upload(self, kind: str, filename: str, data: bytes):
        self._check_ext(kind, filename)
        full = os.path.join(MODELS_DIR, filename)
        with open(full, "wb") as f: f.write(data)
        self.set_active(kind, filename)
        return os.path.basename(full)

    def set_active(self, kind: str, filename: str):
        self._check_ext(kind, filename)
        full = os.path.join(MODELS_DIR, filename)
        if not os.path.exists(full): raise FileNotFoundError(full)
        self.active[kind] = filename
        self._save_meta()
        self.load(kind, full)

    def load(self, kind: str, full: Optional[str] = None):
        if full is None:
            name = self.active.get(kind)
            if not name: self.loaded.pop(kind, None); return None
            full = os.path.join(MODELS_DIR, name)
        obj = joblib.load(full)  # 建议用 sklearn Pipeline 保存
        self.loaded[kind] = ModelPack(kind, full, obj)
        return obj

    # 统一推理：特征顺序需与前端一致
    def predict_crop(self, feats: dict):
        mp = self.loaded.get("crop")
        if not mp: raise RuntimeError("作物推荐模型未加载")
        X = np.array([list(feats.values())], dtype=float)
        y = mp.obj.predict(X)[0]
        # 若是分类且有概率：
        proba = None
        if hasattr(mp.obj, "predict_proba"):
            proba = mp.obj.predict_proba(X)[0].tolist()
        return {"recommended_crop": str(y), "proba": proba}

    def predict_yield(self, feats: dict):
        mp = self.loaded.get("yield")
        if not mp: raise RuntimeError("产量预测模型未加载")
        X = np.array([list(feats.values())], dtype=float)
        val = float(mp.obj.predict(X)[0])
        return {"yield_pred": val, "unit": "kg/mu"}
